msg_tool\scripts\kirikiri/
simple_crypt.rs

1//! Kirikiri Simple Crypt Text File
2use crate::ext::io::*;
3use crate::scripts::base::*;
4use crate::types::*;
5use anyhow::Result;
6use overf::wrapping;
7use std::io::Read;
8
9#[derive(Debug)]
10/// Kirikiri Simple Crypt Script Builder
11pub struct SimpleCryptBuilder {}
12
13impl SimpleCryptBuilder {
14    /// Creates a new instance of `SimpleCryptBuilder`
15    pub fn new() -> Self {
16        Self {}
17    }
18}
19
20impl ScriptBuilder for SimpleCryptBuilder {
21    fn default_encoding(&self) -> Encoding {
22        Encoding::Utf8
23    }
24
25    fn build_script(
26        &self,
27        buf: Vec<u8>,
28        filename: &str,
29        _encoding: Encoding,
30        _archive_encoding: Encoding,
31        _config: &ExtraConfig,
32        _archive: Option<&Box<dyn Script>>,
33    ) -> Result<Box<dyn Script>> {
34        Ok(Box::new(SimpleCrypt::new(buf, filename)?))
35    }
36
37    fn extensions(&self) -> &'static [&'static str] {
38        &[]
39    }
40
41    fn script_type(&self) -> &'static ScriptType {
42        &ScriptType::KirikiriSimpleCrypt
43    }
44
45    fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
46        if buf_len >= 5
47            && buf[0] == 0xfe
48            && buf[1] == 0xfe
49            && (buf[2] == 0 || buf[2] == 1 || buf[2] == 2)
50            && buf[3] == 0xff
51            && buf[4] == 0xfe
52        {
53            Some(10)
54        } else {
55            None
56        }
57    }
58}
59
60#[derive(Debug)]
61/// Kirikiri Simple Crypt Script
62pub struct SimpleCrypt {
63    /// Crypt mode
64    crypt: u8,
65    data: MemReader,
66    ext: String,
67}
68
69impl SimpleCrypt {
70    /// Creates a new `SimpleCrypt` script from the given buffer and filename
71    ///
72    /// * `buf` - The buffer containing the SimpleCrypt data
73    /// * `filename` - The name of the file
74    pub fn new(buf: Vec<u8>, filename: &str) -> Result<Self> {
75        let mut reader = MemReader::new(buf);
76        let mut header = [0u8; 5];
77        reader.read_exact(&mut header)?;
78        if header[0] != 0xfe
79            || header[1] != 0xfe
80            || (header[2] != 0 && header[2] != 1 && header[2] != 2)
81            || header[3] != 0xff
82            || header[4] != 0xfe
83        {
84            return Err(anyhow::anyhow!("Invalid SimpleCrypt header"));
85        }
86        Ok(Self {
87            crypt: header[2],
88            data: reader,
89            ext: std::path::Path::new(filename)
90                .extension()
91                .and_then(|s| s.to_str())
92                .unwrap_or("")
93                .to_string(),
94        })
95    }
96
97    pub(crate) fn unpack(crypt: u8, data: MemReaderRef) -> Result<Vec<u8>> {
98        match crypt {
99            0 => Self::unpack_mode0(data),
100            1 => Self::unpack_mode1(data),
101            2 => Self::unpack_mode2(data),
102            _ => Err(anyhow::anyhow!("Unsupported SimpleCrypt mode: {}", crypt)),
103        }
104    }
105
106    fn unpack_mode0(input: MemReaderRef) -> Result<Vec<u8>> {
107        let mut data = Vec::with_capacity(input.data.len() - 3);
108        data.push(0xff);
109        data.push(0xfe);
110        data.extend_from_slice(&input.data[5..]);
111        for i in 2..data.len() {
112            let ch = data[i] as u16;
113            if ch >= 20 {
114                data[i] = wrapping! {ch ^ (((ch & 0xfe) << 8) ^ 1)} as u8;
115            }
116        }
117        Ok(data)
118    }
119
120    fn unpack_mode1(input: MemReaderRef) -> Result<Vec<u8>> {
121        let mut data = Vec::with_capacity(input.data.len() - 3);
122        data.push(0xff);
123        data.push(0xfe);
124        data.extend_from_slice(&input.data[5..]);
125        for i in 2..data.len() {
126            let mut ch = data[i] as u32;
127            ch = wrapping! {((ch & 0xaaaaaaaa) >> 1) | ((ch & 0x55555555) << 1)};
128            data[i] = ch as u8;
129        }
130        Ok(data)
131    }
132
133    fn unpack_mode2(mut reader: MemReaderRef) -> Result<Vec<u8>> {
134        reader.pos = 5;
135        let compressed = reader.read_u64()?;
136        debug_assert!(compressed + 5 == reader.data.len() as u64);
137        let uncompressed = reader.read_u64()?;
138        let mut stream = flate2::Decompress::new(false);
139        let mut data = Vec::with_capacity(uncompressed as usize + 2);
140        data.push(0xff);
141        data.push(0xfe);
142        data.resize(uncompressed as usize + 2, 0);
143        stream.decompress(
144            &reader.data[reader.pos..],
145            &mut data[2..],
146            flate2::FlushDecompress::Finish,
147        )?;
148        Ok(data)
149    }
150}
151
152impl Script for SimpleCrypt {
153    fn default_output_script_type(&self) -> OutputScriptType {
154        OutputScriptType::Custom
155    }
156
157    fn default_format_type(&self) -> FormatOptions {
158        FormatOptions::None
159    }
160
161    fn is_output_supported(&self, output: OutputScriptType) -> bool {
162        matches!(output, OutputScriptType::Custom)
163    }
164
165    fn custom_output_extension<'a>(&'a self) -> &'a str {
166        &self.ext
167    }
168
169    fn custom_export(&self, filename: &std::path::Path, _encoding: Encoding) -> Result<()> {
170        let data = Self::unpack(self.crypt, self.data.to_ref())?;
171        let mut writer = crate::utils::files::write_file(filename)?;
172        writer.write_all(&data)?;
173        Ok(())
174    }
175}